/*
 * Description: app filter process manager
 */

#include "app_filter_manager.h"

#include <sstream>
#include <string>
#include <thread>
#include <io.h>
#include <vector>
#include <setupapi.h>
#include <TlHelp32.h>
#include "log.h"
#include "../common/guest_utils.h"
#include "../common/session_utils.h"
#include "../common/guest_utils.h"

using namespace COMMON;

namespace {
const std::string FILTER_AGENT_PATH = "C:\\ProgramData\\AppFilter\\FilterAgent.exe";
const std::string DEFAULT_CONFIG_PATH = "C:\\ProgramData\\AppFilter\\config.json";
const std::string PROCESS_NAME = "FilterAgent.exe";
const int TIME_OUT = 3000;
}  // namespace

namespace GUEST {
void AppFilterManager::Init()
{
    SessionUtils::RegSessionChangeCallback("app_filter_agent_", [this](int32_t state) {
        if ((state == WTS_SESSION_LOGON || state == WTS_SESSION_UNLOCK) &&
            watchedSession != WTSGetActiveConsoleSessionId()) {
            StopAll();
        }
        Start();
    });
    Start();
}

void AppFilterManager::Deinit()
{
    SessionUtils::UnRegSessionChangeCallback("app_filter_agent_");
    Stop();
}

void AppFilterManager::Start()
{
    if (isWatched || !SessionUtils::IsUserLogon()) {
        return;
    }
    std::thread t([this]() {
        STARTUPINFO si;
        PROCESS_INFORMATION pi;
        std::stringstream cmd;
        cmd << FILTER_AGENT_PATH << " " << DEFAULT_CONFIG_PATH;

        isWatched = true;
        watchedSession = WTSGetActiveConsoleSessionId();
        while (isWatched) {
            if (!IsReady()) {
                isWatched = false;
                return;
            }
            DWORD pid = 0;
            bool isRunning = GetAgentProcessPID(pid);
            HANDLE process = nullptr;
            if (isRunning && (process = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid)) != nullptr) {
                WaitForSingleObject(process, INFINITE);
                CloseHandle(process);
                continue;
            }
            ZeroMemory(&si, sizeof(si));
            si.cb = sizeof(si);
            ZeroMemory(&pi, sizeof(pi));
            LOG_INFO("start app filter agent process.");
            if (!GuestUtils::CreateUserProcess(cmd.str(), si, pi)) {
                LOG_ERR("pull filter agent process failed.");
                break;
            }
            WaitForSingleObject(pi.hProcess, INFINITE);
            CloseHandle(pi.hProcess);
            CloseHandle(pi.hThread);
            process = nullptr;
        }
    });
    t.detach();
}

void AppFilterManager::Stop()
{
    isWatched = false;
}

bool AppFilterManager::IsReady()
{
    if (_access(FILTER_AGENT_PATH.c_str(), 0) == -1) {
        LOG_ERR("filter_agent is not exist.");
        return false;
    }
    if (_access(DEFAULT_CONFIG_PATH.c_str(), 0) == -1) {
        LOG_ERR("filter_agent config is not exist.");
        return false;
    }
    if (GuestUtils::GetServiceState("app_filter_driver") != SERVICE_RUNNING) {
        LOG_ERR("driver is not ready");
        return false;
    }
    return true;
}

bool AppFilterManager::GetAgentProcessPID(DWORD &pid)
{
    HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (INVALID_HANDLE_VALUE == hSnapshot) {
        return false;
    }
    PROCESSENTRY32 pe = {sizeof(pe)};
    for (BOOL ret = Process32First(hSnapshot, &pe); ret; ret = Process32Next(hSnapshot, &pe)) {
        if (strcmp(pe.szExeFile, PROCESS_NAME.c_str()) != 0) {
            continue;
        }
        DWORD psid = 0;
        if (ProcessIdToSessionId(pe.th32ProcessID, &psid) && psid == WTSGetActiveConsoleSessionId()) {
            pid = pe.th32ProcessID;
            CloseHandle(hSnapshot);
            return true;
        }
    }
    CloseHandle(hSnapshot);
    return false;
}

void AppFilterManager::StopAll()
{
    isWatched = false;
    LOG_INFO("kill app filter agent ...");
    std::string cmd = "taskkill /F /IM FilterAgent.exe";
    if (!GuestUtils::ExecCmd(cmd, false, TIME_OUT)) {
        LOG_ERR("Failed to execute {}", cmd);
    }
}

}  // namespace GUEST